#!/usr/bin/env python
import os, sys
import numpy as np
from astropy.time import Time
import matplotlib.pylab as plt

def read_switchtime(fname):
    st = []
    mjd = []
    wave = []
    fp = open(fname)
    lines = fp.readlines()
    for i in range(len(lines)):
        # print('lines[{}]: {}'.format(i, lines[i]))
        t, w = lines[i][0:15], lines[i][16:]
        year = t[0:4]
        mon  = t[4:6]
        day  = t[6:8]
        hour = t[9:11]
        mint = t[11:13]
        sec  = t[13:15]
        d_iso= year+'-'+mon+'-'+day+'T'+hour+':'+mint+':'+sec
        st.append(d_iso)
        mjd.append(Time(d_iso).mjd)
        wave.append(float(w))
    fp.close()
    
    st = np.array(st)
    mjd = np.array(mjd)
    wave = np.array(wave)
    dT = np.mean(mjd[1:] - mjd[:-1]) / 2
    
    print('dT: {:10.5f} days'.format(dT))
    print('dT: {:10.5f} seconds'.format(dT*86400))
    
    return st, mjd, wave, dT


def read_flux(fname):
    mjd = []
    flux_raw = []
    flux_log = []
    fp = open(fname)

    lines = fp.readlines()

    for i in range(len(lines)):
        d, f = lines[i].split(' ')
        mjd.append(Time(d).mjd)
        f = float(f)
        flux_raw.append(f)
        flux_log.append(np.log10(np.abs(f)))

    fp.close()
    return np.array(mjd), np.array(flux_raw), np.array(flux_log)

def extract_calib_flux(mdir, dmjd=5./86400):

    flux_mjd, flux_raw, flux_log = read_flux(mdir+'/monitor_calib.txt')
    st, st_mjd, wave, dT = read_switchtime(mdir+'/switchtime.txt')

    flux_all = []
    flux_bkg = []
    flux_net = []
    
    st_mjd_a = []
    st_mjd_b = []
    st_mjd_c = []
    st_mjd_d = []
    for i in range(len(st_mjd)):
        idx_all = np.logical_and(   flux_mjd>(st_mjd[i]-2*dT+dmjd), \
                                    flux_mjd<(st_mjd[i]-1*dT-dmjd))

        idx_bkg = np.logical_and(   flux_mjd>(st_mjd[i]-1*dT+dmjd), \
                                    flux_mjd<(st_mjd[i]-0*dT-dmjd))

        st_mjd_a.append(st_mjd[i]-2*dT+dmjd)
        st_mjd_b.append(st_mjd[i]-1*dT-dmjd)
        st_mjd_c.append(st_mjd[i]-1*dT+dmjd)
        st_mjd_d.append(st_mjd[i]-0*dT-dmjd)

        fa = np.median(flux_raw[idx_all])
        fb = np.median(flux_raw[idx_bkg])
        flux_all.append(fa)
        flux_bkg.append(fb)
        flux_net.append(fa-fb)
        
    plt.figure()
    plt.plot(flux_mjd, flux_raw, '.')
    plt.vlines(st_mjd, ymin=np.min(flux_raw), ymax=np.max(flux_raw), linestyle='dashed')
    plt.vlines(st_mjd_a, ymin=np.min(flux_raw), ymax=np.max(flux_raw), linestyle='dashed', color='b')
    plt.vlines(st_mjd_b, ymin=np.min(flux_raw), ymax=np.max(flux_raw), linestyle='dashed', color='b')
    plt.vlines(st_mjd_c, ymin=np.min(flux_raw), ymax=np.max(flux_raw), linestyle='dashed', color='r')
    plt.vlines(st_mjd_d, ymin=np.min(flux_raw), ymax=np.max(flux_raw), linestyle='dashed', color='r')
    plt.show()
    
    plt.figure()
    flux_bkg = np.array(flux_bkg)
    plt.plot((flux_bkg-flux_bkg.mean())/flux_bkg.mean(), label='background')
    plt.show()

    return np.array(flux_net), np.array(wave)

def extract_sphere_flux(mdir, dmjd=5./86400):

    flux_mjd, flux_raw, flux_log = read_flux(mdir+'/monitor_sphere.txt')
    st, st_mjd, wave, dT = read_switchtime(mdir+'/switchtime.txt')

    flux_all = []
    flux_bkg = []
    flux_net = []
    for i in range(len(st_mjd)):
        idx_all = np.logical_and(   flux_mjd>(st_mjd[i]-2*dT+dmjd), \
                                    flux_mjd<(st_mjd[i]-1*dT-dmjd))

        idx_bkg = np.logical_and(   flux_mjd>(st_mjd[i]-1*dT+dmjd), \
                                    flux_mjd<(st_mjd[i]-0*dT-dmjd))

        fa = np.median(np.abs(flux_raw[idx_all]))
        fb = np.median(np.abs(flux_raw[idx_bkg]))
        flux_all.append(fa)
        flux_bkg.append(fb)
        flux_net.append( np.abs(fa-fb) )

    return np.array(flux_net), np.array(wave)


def main():
    """
    TODO: 增加依据积分球监测数据做修正的功能[done]
    """
    if len(sys.argv) != 4:
        print('usage: {} dir_no_win  dir_win output_filename'.format(sys.argv[0]))
        sys.exit(0)

    dir_no_win = sys.argv[1] # data without glass
    dir_win = sys.argv[2] # data with glass

    flux_calib1, wave_calib1 = extract_calib_flux(dir_no_win)
    flux_sphere1, wave_sphere1 = extract_sphere_flux(dir_no_win)

    flux_calib2, wave_calib2 = extract_calib_flux(dir_win)
    flux_sphere2, wave_sphere2 = extract_sphere_flux(dir_win)
    
    trans = (flux_calib2/flux_sphere2*flux_sphere1)/flux_calib1
    
    plt.figure()
    plt.plot(wave_calib1, flux_calib2/flux_calib1, '--o', label='no sphere correction')
    plt.plot(wave_calib1, trans, '--s', label='after sphere correction')
    # plt.plot(wave_calib1, flux_sphere2/flux_sphere1, '--o')

    fp = open(sys.argv[3],'w')
    for i in range(len(trans)):
        fp.write('{:.10f} {:.10f}\n'.format(wave_calib1[i], trans[i]))
    fp.close()    

    plt.legend(loc='best')
    plt.show()


if __name__ == '__main__':
    main()

